#include <functional>
#include <random>

namespace qing {
    class DataGenerator {
    private:
        std::random_device rd;
        std::mt19937 gen;
        
    public:
        DataGenerator() : gen(rd()) {}
        
        // 生成线性关系数据：y = w*x + b + noise
        void generate_linear_data(int num_samples, int input_dim, 
                                std::vector<std::vector<double>>& features,
                                std::vector<std::vector<double>>& labels,
                                double noise_std = 0.1) {
            std::normal_distribution<double> noise_dist(0.0, noise_std);
            std::normal_distribution<double> weight_dist(0.0, 1.0);
            
            // 生成真实的权重和偏置
            std::vector<double> true_weights(input_dim);
            double true_bias = weight_dist(gen);
            
            for (int i = 0; i < input_dim; ++i) {
                true_weights[i] = weight_dist(gen);
            }
            
            features.clear();
            labels.clear();
            
            for (int i = 0; i < num_samples; ++i) {
                std::vector<double> sample(input_dim);
                double y = true_bias;
                
                for (int j = 0; j < input_dim; ++j) {
                    sample[j] = weight_dist(gen);  // 随机生成输入特征
                    y += true_weights[j] * sample[j];
                }
                
                // 添加噪声
                y += noise_dist(gen);
                
                features.push_back(sample);
                labels.push_back({y});
            }
            
            std::cout << "Generated linear data with " << num_samples << " samples" << std::endl;
            std::cout << "True function: y = ";
            for (int i = 0; i < input_dim; ++i) {
                std::cout << true_weights[i] << "*x" << i;
                if (i < input_dim - 1) std::cout << " + ";
            }
            std::cout << " + " << true_bias << std::endl;
        }
        
        // 生成二分类数据
        void generate_binary_classification_data(int num_samples, int input_dim,
                                            std::vector<std::vector<double>>& features,
                                            std::vector<std::vector<double>>& labels) {
            std::normal_distribution<double> feature_dist(0.0, 1.0);
            std::uniform_real_distribution<double> weight_dist(-1.0, 1.0);
            
            // 生成真实的权重和偏置
            std::vector<double> true_weights(input_dim);
            double true_bias = weight_dist(gen);
            
            for (int i = 0; i < input_dim; ++i) {
                true_weights[i] = weight_dist(gen);
            }
            
            features.clear();
            labels.clear();
            
            for (int i = 0; i < num_samples; ++i) {
                std::vector<double> sample(input_dim);
                double score = true_bias;
                
                for (int j = 0; j < input_dim; ++j) {
                    sample[j] = feature_dist(gen);
                    score += true_weights[j] * sample[j];
                }
                
                // 使用sigmoid函数生成概率，然后二值化
                double probability = 1.0 / (1.0 + std::exp(-score));
                int label = (probability > 0.5) ? 1 : 0;
                
                features.push_back(sample);
                labels.push_back({static_cast<double>(label)});
            }
            
            std::cout << "Generated binary classification data with " << num_samples << " samples" << std::endl;
        }
        
        // 生成异或问题数据（非线性可分）
        void generate_xor_data(int num_samples_per_class,
                            std::vector<std::vector<double>>& features,
                            std::vector<std::vector<double>>& labels) {
            std::normal_distribution<double> noise_dist(0.0, 0.1);
            
            features.clear();
            labels.clear();
            
            // 类别1: (0,0) 和 (1,1)
            for (int i = 0; i < num_samples_per_class; ++i) {
                features.push_back({0.0 + noise_dist(gen), 0.0 + noise_dist(gen)});
                labels.push_back({0.0});
                
                features.push_back({1.0 + noise_dist(gen), 1.0 + noise_dist(gen)});
                labels.push_back({0.0});
            }
            
            // 类别2: (0,1) 和 (1,0)
            for (int i = 0; i < num_samples_per_class; ++i) {
                features.push_back({0.0 + noise_dist(gen), 1.0 + noise_dist(gen)});
                labels.push_back({1.0});
                
                features.push_back({1.0 + noise_dist(gen), 0.0 + noise_dist(gen)});
                labels.push_back({1.0});
            }
            
            std::cout << "Generated XOR data with " << features.size() << " samples" << std::endl;
        }
    };
}